"""This module contains objective functions for the model.
"""

import pyoptinterface as poi
from loguru import logger

class AddObjective:
    """Objective function class to determine the total spillflow of the model.
    """
    def __init__(self, model : object) -> None:
        """The constructor for objective functions class.

        Parameters
        ----------
        model : object
            Model to be solved.
        """
        self.model = model
        self.define_objective()

    # def define_objective(self) -> None:
    #     """Objective function of the model.
    #     """
    #     model = self.model
    #     unit_name = "葠窝水库"
    #     model.end_storage_flood_season = self.unit_end_storage_flood_season_rule(unit_name)
    #     model.set_objective(model.end_storage_flood_season, sense=poi.ObjectiveSense.Maximize)

    # def define_objective(self) -> None:
    #     """Objective function of the model.
    #     """
    #     model = self.model
    #     model.end_storage_flood_season = self.system_end_storage_flood_season_rule()
    #     model.set_objective(model.end_storage_flood_season, sense=poi.ObjectiveSense.Maximize)

    # def define_objective(self) -> None:
    #     """Objective function of the model.
    #     """
    #     model = self.model
    #     model.spillflow_total = self.system_spillflow_calc_rule()
    #     model.set_objective(model.spillflow_total, sense=poi.ObjectiveSense.Minimize)


    def define_objective(self) -> None:
        """Objective function of the model.
        根据 obj_cons.json 中的配置，支持多目标、多种目标类型（spillflow、end_storage_flood_season）。
        支持混合优化方向（minimize和maximize），自动统一转换为minimize。
        当混合不同类型目标时，自动进行归一化处理。
        """
        model = self.model
        objectives = model.param['objectives']
        
        # 检查是否有混合优化方向
        self._check_mixed_optimization_sense(objectives)
        
        # 检查是否需要归一化（混合不同类型的目标）
        needs_normalization = self._check_needs_normalization(objectives)
        
        # 构建目标项（会自动处理优化方向和归一化）
        weighted_terms, objective_terms_by_type = self._build_objective_terms(
            objectives, needs_normalization
        )
        
        # 统一使用 minimize 方向
        objective_expr = poi.quicksum(weighted_terms)
        model.set_objective(objective_expr, sense=poi.ObjectiveSense.Minimize)
        
        # 缓存各类目标表达式供结果输出使用
        self._cache_objective_expressions(objective_terms_by_type)
        
        logger.info(f"目标函数设置完成，统一优化方向：minimize")

    def _check_mixed_optimization_sense(self, objectives: list) -> None:
        """检查是否有混合优化方向，并记录日志。
        
        Parameters
        ----------
        objectives : list
            目标函数配置列表
        """
        senses = set()
        for obj in objectives:
            sense = obj.get('sense', '').strip().lower()
            senses.add(sense)
        
        if len(senses) > 1:
            logger.info(
                f"检测到混合优化方向: {', '.join(senses)}。"
                "系统将统一转换为 minimize，maximize 目标将添加负号。"
            )
        else:
            sense = list(senses)[0] if senses else 'minimize'
            logger.info(f"所有目标使用统一优化方向: {sense}")
    
    def _check_needs_normalization(self, objectives: list) -> bool:
        """检查是否需要归一化处理（当有多个目标时）。
        
        Parameters
        ----------
        objectives : list
            目标函数配置列表
            
        Returns
        -------
        bool
            True 表示需要归一化（多目标），False 表示不需要（单目标）
        """
        needs_norm = len(objectives) > 1
        
        if needs_norm:
            logger.info(
                f"检测到多目标优化（{len(objectives)}个目标）。"
                "系统将对目标进行归一化处理以避免量纲问题。"
            )
        else:
            logger.info(f"单目标优化，无需归一化处理")
        
        return needs_norm
    
    def _build_objective_terms(self, objectives: list, needs_normalization: bool = False) -> tuple:
        """构建目标项，自动处理混合优化方向和归一化。
        
        所有目标统一转换为 minimize 方向：
        - minimize 目标：保持原样
        - maximize 目标：添加负号（minimize -f 等价于 maximize f）
        
        当混合不同类型目标时，进行归一化处理：
        - spillflow: 除以总弃水上限
        - end_storage: 除以总允许库容上限
        
        Parameters
        ----------
        objectives : list
            目标函数配置列表
        needs_normalization : bool
            是否需要归一化处理（混合类型目标时为 True）
            
        Returns
        -------
        tuple
            (weighted_terms, objective_terms_by_type)
            - weighted_terms: 所有加权后并调整方向的目标项列表
            - objective_terms_by_type: 按类型分组的目标项字典
        """
        weighted_terms = []
        objective_terms_by_type = {
            'spillflow': [],
            'end_storage_flood_season': [],
            'end_storage_dispatch_term': []
        }
        
        for obj in objectives:
            obj_name = obj.get('name', '').strip().lower()
            scope = obj.get('scope', 'system').strip().lower()
            unit_name = obj.get('unit_stcd', '')
            weight = self._parse_weight(obj.get('weight', 1.0))
            sense = obj.get('sense', 'minimize').strip().lower()
            
            # 根据目标类型构建表达式
            term_expr = self._build_single_objective_term(
                obj_name, scope, unit_name, weight, needs_normalization
            )
            
            if term_expr is not None:
                # 应用权重
                weighted_term = term_expr * weight
                
                # 处理优化方向：统一转换为 minimize
                # maximize f(x) 等价于 minimize -f(x)
                if sense == 'maximize':
                    weighted_term = -1 * weighted_term  # 使用乘法而不是一元负号
                    logger.debug(f"目标 {obj_name} 为 maximize，已转换为 minimize -f")
                
                weighted_terms.append(weighted_term)
                
                # 按类型分组收集（保存原始带权重和方向调整后的项）
                if 'spillflow' in obj_name:
                    objective_terms_by_type['spillflow'].append(weighted_term)
                elif 'end_storage_flood_season' in obj_name:
                    objective_terms_by_type['end_storage_flood_season'].append(weighted_term)
                elif 'end_storage_dispatch_term' in obj_name:
                    objective_terms_by_type['end_storage_dispatch_term'].append(weighted_term)
        
        if not weighted_terms:
            logger.error("没有有效的目标函数，无法构建模型")
            raise ValueError("未定义有效的目标函数")
        
        return weighted_terms, objective_terms_by_type
    
    def _parse_weight(self, weight_val) -> float:
        """解析权重值。
        
        Parameters
        ----------
        weight_val : any
            权重配置值
            
        Returns
        -------
        float
            解析后的权重，默认为 1.0
        """
        try:
            return float(weight_val) if weight_val is not None else 1.0
        except (TypeError, ValueError):
            logger.warning(f"无效的权重值 {weight_val}，使用默认值 1.0")
            return 1.0
    
    def _build_single_objective_term(self, obj_name: str, scope: str, unit_name: str, 
                                     weight: float, needs_normalization: bool = False):
        """构建单个目标项。
        
        Parameters
        ----------
        obj_name : str
            目标名称
        scope : str
            作用域 ('system' 或 'unit')
        unit_name : str
            单元名称（当 scope='unit' 时需要）
        weight : float
            权重
        needs_normalization : bool
            是否需要归一化处理
            
        Returns
        -------
        poi.ExprBuilder or None
            目标表达式，如果无法构建则返回 None
        """
        if 'spillflow' in obj_name:
            return self._build_spillflow_term(scope, unit_name, weight, needs_normalization)
        elif 'end_storage_flood_season' in obj_name:
            return self._build_end_storage_term(scope, unit_name, weight, needs_normalization)
        elif 'end_storage_dispatch_term' in obj_name:
            return self._build_end_storage_dispatch_term(scope, unit_name, weight, needs_normalization)
        else:
            logger.warning(f"未识别的目标类型：{obj_name}，跳过")
            return None
    
    def _build_spillflow_term(self, scope: str, unit_name: str, weight: float, 
                             needs_normalization: bool = False):
        """构建弃水目标项。
        
        Parameters
        ----------
        scope : str
            作用域
        unit_name : str
            单元名称
        weight : float
            权重
        needs_normalization : bool
            是否需要归一化（混合目标类型时进行归一化）
        """
        if scope == 'unit':
            if not unit_name:
                logger.error(f"弃水目标为 unit 级别，但未指定 unit_stcd")
                return None
            logger.info(f"添加目标：{unit_name} 的弃水量，权重={weight}")
            term = self.unit_spillflow_calc_rule(unit_name)
        else:
            logger.info(f"添加目标：系统总弃水量，权重={weight}")
            term = self.system_spillflow_calc_rule()
        
        # 如果需要归一化（混合目标类型），则除以弃水上限总和
        if needs_normalization:
            normalization_factor = self._get_spillflow_normalization_factor(scope, unit_name)
            if normalization_factor > 0:
                term = term / normalization_factor
                logger.info(f"弃水目标已归一化，归一化因子={normalization_factor:.2f}")
        
        return term
    
    def _build_end_storage_term(self, scope: str, unit_name: str, weight: float,
                                needs_normalization: bool = False):
        """构建汛末蓄水目标项。
        
        Parameters
        ----------
        scope : str
            作用域
        unit_name : str
            单元名称
        weight : float
            权重
        needs_normalization : bool
            是否需要归一化（混合目标类型时归一化）
        """
        if scope == 'unit':
            if not unit_name:
                logger.error(f"汛末蓄水目标为 unit 级别，但未指定 unit_stcd")
                return None
            logger.info(f"添加目标：{unit_name} 的汛末蓄水库容，权重={weight}")
            term = self.unit_end_storage_flood_season_rule(unit_name)
        else:
            logger.info(f"添加目标：系统汛末蓄水库容，权重={weight}")
            term = self.system_end_storage_flood_season_rule()
        
        # 如果需要归一化（混合目标类型），则除以库容范围总和
        if needs_normalization:
            normalization_factor = self._get_end_storage_flood_season_normalization_factor(scope, unit_name)
            if normalization_factor > 0:
                term = term / normalization_factor
                logger.info(f"蓄水目标已归一化，归一化因子={normalization_factor:.2f}")
        
        return term
    
    def _build_end_storage_dispatch_term(self, scope: str, unit_name: str, weight: float,
                                        needs_normalization: bool = False):
        """构建调度期末库容目标项。
        
        Parameters
        ----------
        scope : str
            作用域
        unit_name : str
            单元名称
        weight : float
            权重
        needs_normalization : bool
            是否需要归一化（混合目标类型时归一化）
        """
        if scope == 'unit':
            if not unit_name:
                logger.error(f"调度期末库容目标为 unit 级别，但未指定 unit_stcd")
                return None
            logger.info(f"添加目标：{unit_name} 的调度期末库容，权重={weight}")
            term = self.unit_end_storage_dispatch_term_rule(unit_name)
        else:
            logger.info(f"添加目标：系统调度期末库容，权重={weight}")
            term = self.system_end_storage_dispatch_term_rule()
        
        # 如果需要归一化（混合目标类型），则除以库容范围总和
        if needs_normalization:
            normalization_factor = self._get_end_storage_dispatch_term_normalization_factor(scope, unit_name)
            if normalization_factor > 0:
                term = term / normalization_factor
                logger.info(f"调度期末库容目标已归一化，归一化因子={normalization_factor:.2f}")
        
        return term
    
    def _get_spillflow_normalization_factor(self, scope: str, unit_name: str = None) -> float:
        """获取弃水目标的归一化因子（弃水上限总和）。
        
        Parameters
        ----------
        scope : str
            作用域
        unit_name : str, optional
            单元名称（当 scope='unit' 时使用）
            
        Returns
        -------
        float
            归一化因子（弃水上限总和）
        """
        model = self.model
        spillflow_up = model.param.get('spillflow_up', {})
        
        if not spillflow_up:
            logger.warning("未配置 spillflow_up，无法归一化弃水目标")
            return 1.0
        
        total = 0.0
        if scope == 'unit' and unit_name:
            # Unit 级别：只计算该水库的弃水上限总和
            for (r, t), value in spillflow_up.items():
                if r == unit_name:
                    total += value
        else:
            # System 级别：计算所有水库的弃水上限总和
            for value in spillflow_up.values():
                total += value
        
        return total if total > 0 else 1.0
    
    def _get_end_storage_flood_season_normalization_factor(self, scope: str, unit_name: str = None) -> float:
        """获取蓄水目标的归一化因子（库容范围总和）。
        
        归一化因子 = Σ Smax，即所有水库的库容范围之和
        
        Parameters
        ----------
        scope : str
            作用域
        unit_name : str, optional
            单元名称（当 scope='unit' 时使用）
            
        Returns
        -------
        float
            归一化因子（库容范围总和）
        """
        model = self.model
        storage_up = model.param.get('storage_up', {})
        
        if not storage_up:
            logger.warning("未配置 storage_up，无法归一化蓄水目标")
            return 1.0
        
        # 获取汛末时间
        flood_end_time = self._get_flood_season_end_time()
        
        total = 0.0
        if scope == 'unit' and unit_name:
            # Unit 级别：只计算该水库的库容范围
            if (unit_name, flood_end_time) in storage_up:
                Smax = storage_up[unit_name, flood_end_time]
                total = Smax
        else:
            # System 级别：计算所有水库的库容范围总和
            for r in model.reservoir:
                if (r, flood_end_time) in storage_up:
                    Smax = storage_up[r, flood_end_time]
                    total += Smax
        
        return total if total > 0 else 1.0
    
    def _get_end_storage_dispatch_term_normalization_factor(self, scope: str, unit_name: str = None) -> float:
        """获取调度期末库容目标的归一化因子（库容范围总和）。
        
        归一化因子 = Σ Smax，即所有水库的调度期末库容上限之和
        
        Parameters
        ----------
        scope : str
            作用域
        unit_name : str, optional
            单元名称（当 scope='unit' 时使用）
            
        Returns
        -------
        float
            归一化因子（库容范围总和）
        """
        model = self.model
        storage_up = model.param.get('storage_up', {})
        
        if not storage_up:
            logger.warning("未配置 storage_up，无法归一化调度期末库容目标")
            return 1.0
        
        # 获取调度期末时间（最后一个时间步）
        dispatch_end_time = model.time_p[-1] if hasattr(model, 'time_p') and model.time_p else model.time[-1]
        
        total = 0.0
        if scope == 'unit' and unit_name:
            # Unit 级别：只计算该水库的库容范围
            if (unit_name, dispatch_end_time) in storage_up:
                Smax = storage_up[unit_name, dispatch_end_time]
                total = Smax
        else:
            # System 级别：计算所有水库的库容范围总和
            for r in model.reservoir:
                if (r, dispatch_end_time) in storage_up:
                    Smax = storage_up[r, dispatch_end_time]
                    total += Smax
        
        return total if total > 0 else 1.0
    
    def _cache_objective_expressions(self, objective_terms_by_type: dict) -> None:
        """缓存各类目标的表达式供结果输出使用。
        
        Parameters
        ----------
        objective_terms_by_type : dict
            按类型分组的目标项字典
        """
        model = self.model
        
        # 缓存弃水目标
        spillflow_terms = objective_terms_by_type.get('spillflow', [])
        if spillflow_terms:
            model.spillflow_total = poi.quicksum(spillflow_terms)
            logger.info("已保存 spillflow_total 表达式用于结果输出")
        else:
            model.spillflow_total = 0.0
            logger.info("无弃水目标，spillflow_total 设为 0")
        
        # 缓存汛末蓄水目标
        end_storage_terms = objective_terms_by_type.get('end_storage_flood_season', [])
        if end_storage_terms:
            model.end_storage_flood_season_total = poi.quicksum(end_storage_terms)
            logger.info("已保存 end_storage_flood_season_total 表达式用于结果输出")
        else:
            model.end_storage_flood_season_total = 0.0
            logger.info("无汛末蓄水目标，end_storage_flood_season_total 设为 0")
        
        # 缓存调度期末库容目标
        end_storage_dispatch_terms = objective_terms_by_type.get('end_storage_dispatch_term', [])
        if end_storage_dispatch_terms:
            model.end_storage_dispatch_term_total = poi.quicksum(end_storage_dispatch_terms)
            logger.info("已保存 end_storage_dispatch_term_total 表达式用于结果输出")
        else:
            model.end_storage_dispatch_term_total = 0.0
            logger.info("无调度期末库容目标，end_storage_dispatch_term_total 设为 0")

    def spillflow_calc_breakdown(self, r : str, t : str) -> poi.ExprBuilder:
        """Calculate the spillflow of reservoir r at time t.

        Parameters
        ----------
        r : str
            Reservoir name.
        t : str
            Time.

        Returns
        -------
        poi.ExprBuilder
            Spillflow of reservoir r at time t.
        """
        model = self.model
        return model.spillflow[r, t]

    def end_storage_flood_season_breakdown(self, r : str) -> poi.ExprBuilder:
        """Calculate the water level of reservoir r at the end of the flood season (10月份).

        Parameters
        ----------
        r : str
            Reservoir name.
        
        Returns
        -------
        poi.ExprBuilder
            Storage ratio of reservoir r at the end of flood season.
            (St - Smin) / (Smax - Smin)
        """
        model = self.model
        # 选取汛末（10月份）的时间；若不存在则报错
        flood_end_time = self._get_flood_season_end_time()
        St = model.storage[r, flood_end_time]
        # Smin = model.param['storage_down'][r, flood_end_time]
        # Smax = model.param['storage_up'][r, flood_end_time]
        # return (St - Smin) / (Smax - Smin)
        return St

    def end_storage_dispatch_term_breakdown(self, r : str) -> poi.ExprBuilder:
        """Calculate the water level of reservoir r at the end of the dispatch term.

        Parameters
        ----------
        r : str
            Reservoir name.
        
        Returns
        -------
        poi.ExprBuilder
            Water level of reservoir r at the end of the dispatch term.
        """
        
        model = self.model
        return model.storage[r, model.time_p[-1]]

    def system_spillflow_calc_rule(self) -> poi.ExprBuilder:
        """Calculate the spillflow of the model.

        Returns
        -------
        poi.ExprBuilder
            Spillflow of the model.
        """
        model = self.model
        model.system_spillflow_calc_breakdown = poi.make_tupledict(
            model.reservoir, model.time,
            rule=self.spillflow_calc_breakdown
        )
        return poi.quicksum(model.system_spillflow_calc_breakdown)

    def unit_spillflow_calc_rule(self, unit_name: str) -> poi.ExprBuilder:
        """Calculate the spillflow of a single unit across all times.

        Parameters
        ----------
        unit_name : str
            Unit name.(Such as Reservoir)

        Returns
        -------
        poi.ExprBuilder
            Spillflow of the specified unit across all times.
        """
        model = self.model
        # 将目标单位名称放入模型的集合，使用与系统相同的迭代风格
        model.unit = [unit_name]
        model.unit_spillflow_calc_breakdown = poi.make_tupledict(
            model.unit, model.time,
            rule=self.spillflow_calc_breakdown
        )
        return poi.quicksum(model.unit_spillflow_calc_breakdown)

    def system_end_storage_flood_season_rule(self) -> poi.ExprBuilder:
        """Calculate the water level of all reservoirs at the end of the flood season.

        Returns
        -------
        poi.ExprBuilder
            Water level of the model at the end of the flood season.
         """
        model = self.model
        model.system_end_storage_flood_season_breakdown = poi.make_tupledict(
            model.reservoir,
            rule=self.end_storage_flood_season_breakdown
        )
        return poi.quicksum(model.system_end_storage_flood_season_breakdown)

    def unit_end_storage_flood_season_rule(self, unit_name: str) -> poi.ExprBuilder:
        """Calculate the water level of the specified unit at the end of the flood season.

        Parameters
        ----------
        unit_name : str
            Unit name.(Such as Reservoir)

        Returns
        -------
        poi.ExprBuilder
            Water level of the specified unit at the end of the flood season.
        """
        model = self.model
        model.unit = [unit_name]
        model.unit_end_storage_flood_season_breakdown = poi.make_tupledict(
            model.unit,
            rule=self.end_storage_flood_season_breakdown
        )
        return poi.quicksum(model.unit_end_storage_flood_season_breakdown)

    def system_end_storage_dispatch_term_rule(self) -> poi.ExprBuilder:
        """Calculate the water level of all reservoirs at the end of the dispatch term.

        Returns
        -------
        poi.ExprBuilder
            Water level of the model at the end of the dispatch term.
        """
        model = self.model
        model.system_end_storage_dispatch_term_breakdown = poi.make_tupledict(
            model.reservoir,
            rule=self.end_storage_dispatch_term_breakdown
        )
        return poi.quicksum(model.system_end_storage_dispatch_term_breakdown)

    def unit_end_storage_dispatch_term_rule(self, unit_name: str) -> poi.ExprBuilder:
        """Calculate the water level of the specified unit at the end of the dispatch term.

        Parameters
        ----------
        unit_name : str
            Unit name.(Such as Reservoir)

        Returns
        -------
        poi.ExprBuilder
            Water level of the specified unit at the end of the dispatch term.
        """
        model = self.model
        model.unit = [unit_name]
        model.unit_end_storage_dispatch_term_breakdown = poi.make_tupledict(
            model.unit,
            rule=self.end_storage_dispatch_term_breakdown
        )
        return poi.quicksum(model.unit_end_storage_dispatch_term_breakdown)

    def _get_flood_season_end_time(self) -> str:
        """Return the time index for flood season end (October).

        If October does not exist in `model.time`, log an error and
        fall back to the last available time index.

        Returns
        -------
        str
            A time string like 'YYYY-10'. If missing, the last item of `model.time`.
        """
        model = self.model
        # 从时间序列中筛选10月份（格式 YYYY-10）
        october_times = [t for t in model.time if isinstance(t, str) and t.endswith('-10')]
        if not october_times:
            logger.error("不存在汛末时间（10月份）")
        # 选择最新的一个10月（如果跨年，则取年份最大的10月）
        return sorted(october_times)[-1]